#!/usr/bin/env python
# coding: utf8

from math import log

class BinomialObj(object):
    def __init__(self, numerators, denomenators):
        self.numerators = numerators
        self.denomenators = denomenators
    
    def copy(self,):
        bin_obj = BinomialObj(self.numerators[:], self.denomenators[:])
        return bin_obj
        
    def __mul__(self, multiplier):
        bin_obj = self.copy()
        if isinstance(multiplier, BinomialObj):
            bin_obj.numerators.extend(multiplier.numerators)
            bin_obj.denomenators.extend(multiplier.denomenators)
        elif multiplier is int:
            bin_obj.numerators.append(multiplier)
        else:
            print type(multiplier)
            raise Exception('Unknown type of multiplier')
        return bin_obj
        
    def __div__(self, divisor):
        bin_obj = self.copy()
        if isinstance(divisor, BinomialObj):
            bin_obj.numerators.extend(divisor.denomenators)
            bin_obj.denomenators.extend(divisor.numerators)
        elif divisor is int:
            bin_obj.denomenators.append(divisor)
        else:
            raise Exception('Unknown type of divisor')
        return bin_obj
        
    def _numerator(self,):
        return reduce(lambda x,y: x*y, self.numerators, 1.0)
        
    def _denomenator(self,):
        return reduce(lambda x,y: x*y, self.denomenators, 1.0)
            
    def calculate(self):
        common = set(self.numerators).intersection(self.denomenators)
        while common:
            for number in common:
                self.numerators.remove(number)
                self.denomenators.remove(number)
            common = set(self.numerators).intersection(self.denomenators)
        return float(self._numerator()) / self._denomenator()

class BinomialCoef(BinomialObj):
    def __init__(self, n, k):
        super(BinomialCoef, self).__init__(range(n-k+1, n+1), range(2,k+1))

def theor_probability(l, n1, n2, n12):
    bin_obj = BinomialCoef(n1, n12) * BinomialCoef(l-n1, n2-n12) / BinomialCoef(l, n2)
    return bin_obj.calculate()

def mutual_information(l, n1, n2, n12):
    """
    Взаимная информация между двумя последовательностями длины l
    n1 - число "1" в 1й последовательности
    n2 - число "1" в 2й последовательности
    n12 - число "1" в пересечении последовательностей
    """
    l = float(l)
    n12_th = n1 * n2 / float(l)
    n0_th  = float(n1 - n12_th + n2 - n12_th)
    n1_th  = l - n0_th
    p0_th = n0_th / l
    p1_th = n1_th / l
    n0_exp = float(n1 - n12 + n2 - n12)
    n1_exp = l - n0_exp
    #print 'input: ', l, n1, n2, n12
    #print 'th: ', n0_th, n1_th
    #print 'exp:', n0_exp, n1_exp
    if (n0_exp > 0) and (n1_exp > 0) and (n0_th) > 0 and (n1_th > 0):
        I = n1_exp * log(n1_exp) + n0_exp * log(n0_exp) \
            - n1_exp * log(p1_th) - n0_exp * (p0_th and log(p0_th) or 1.0) - l * log(l)
    else:
        I = 0.0
    return I

def gene_distance(length, set1, set2):
    """
    Расстояние между двумя генами,
    представленными сетами id геномов, в которых они были найдены
    length - размер множества геномов
    """
    return mutual_information(length, len(set1), len(set2), len(set1.intersection(set2)))

def _monte_carlo(length, s1_length, s2_length, num):
    import random
    #random.seed(random.random())
    def random_set(l,n):
        s = set()
        while len(s) < n:
            s.add(random.randrange(l))
        return s

    for i in xrange(num):
        yield gene_distance(length, random_set(length, s1_length), random_set(length, s2_length))


if __name__ == "__main__":

    print gene_distance(1200, set([1,2,3,4]), set([1,2,3,4]))
    print gene_distance(1200, set([1,2,3,4]), set([1,2,3,4,5]))
    print gene_distance(1200, set([1,2,3,4]), set([1,2,3,7,8]))
    print gene_distance(1200, set([1,2,3,4]), set([1,2,7,8,9]))
    print gene_distance(1200, set([1,2,3,4]), set([1,5,6,7,8,9]))
    print gene_distance(1200, set([1,2,3,4]), set([6,7,8,9]))

    print monte_carlo(1200, 4, 5, 1000000, 1.35)

